import os
import pytesseract
from pdf2image import convert_from_path
import re
import pandas as pd
import json
from PIL import Image, ImageEnhance, ImageFilter

# 加载行政区划代码的 JSON 文件
with open('regions.json', 'r', encoding='utf-8') as file:
    region_codes = json.load(file)


# 图像预处理
# def preprocess_image(image):
#     # 转为灰度
#     image = image.convert('L')
#     # 增加对比度
#     enhancer = ImageEnhance.Contrast(image)
#     image = enhancer.enhance(2)
#     # 二值化
#     image = image.point(lambda x: 0 if x < 140 else 255)
#     # 放大图像
#     image = image.resize((int(image.width * 1.5), int(image.height * 1.5)), Image.LANCZOS)
#     return image


# 提取图片中的文本
def extract_text_from_image(image):
    # preprocessed_image = preprocess_image(image)
    return pytesseract.image_to_string(image, lang='chi_sim')  # 使用简体中文 OCR


# 从文本中提取发票信息
def extract_invoice_info(text, file_name):
    invoice_code_match = re.search(r'(\d{10})', text)
    invoice_code = invoice_code_match.group(1) if invoice_code_match else None

    # 从文件名中提取发票号码
    invoice_number = os.path.splitext(file_name)[0]  # 假设文件名即为发票号码

    # 使用更宽松的正则表达式匹配日期
    invoice_date_match = re.search(r'(\d{4})\s*[年./-]\s*(\d{1,2})\s*[月./-]\s*(\d{1,2})', text)
    invoice_date = (
        f"{invoice_date_match.group(1)}/{invoice_date_match.group(2).zfill(2)}/{invoice_date_match.group(3).zfill(2)}"
        if invoice_date_match else None
    )

    invoice_amount = re.search(r'合计[:：]?\s*¥?(\d+\.\d+)', text)
    invoice_type = '增值税专用发票'  # 发票类型固定

    # 根据发票代码前四位来查找对应的行政区划名称
    region_name = region_codes.get(invoice_code[:4], "未知地区") if invoice_code else "未知地区"

    return (
        invoice_code,
        invoice_number,
        invoice_date,
        invoice_type,
        invoice_amount.group(1) if invoice_amount else None,
        # region_name
    )


# 处理 PDF，转换为图片并提取信息
def process_pdf(pdf_path, file_name):
    images = convert_from_path(pdf_path)
    for image in images:
        text = extract_text_from_image(image)
        # print(text)
        return extract_invoice_info(text, file_name)


# 遍历文件夹中的所有PDF文件
def traverse_folders(folder_path):
    data = []
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.pdf'):
                pdf_path = os.path.join(root, file)
                invoice_data = process_pdf(pdf_path, file)
                data.append(invoice_data)
    return data


# 主函数：提取所有发票信息并保存为 Excel 文件
def main():
    folder_path = input("请输入一级文件夹路径: ")
    invoice_data = traverse_folders(folder_path)
    df = pd.DataFrame(invoice_data,
                      columns=['发票代码', '发票号码', '开票日期yyyy/MM/dd', '发票类型', '发票金额（不含税）'])
    # print(df)
    output_path = os.path.join(folder_path, '税票信息.xlsx')    # 保存文件到指定的一级文件夹下
    df.to_excel(output_path, index=False)
    print(f"税票信息已保存至 {output_path}！")

if __name__ == "__main__":
    main()
